[SPARK-48290][CORE] Record skewed shuffle block sizes by default so that AQE can detect skew - #57528
[SPARK-48290][CORE] Record skewed shuffle block sizes by default so that AQE can detect skew#57528Vivek1106-04 wants to merge 2 commits into
Conversation
…hat AQE can detect skew Once a shuffle has more than spark.shuffle.minNumPartitionsToHighlyCompress (2000) partitions, map statuses switch to HighlyCompressedMapStatus, which reports every block below spark.shuffle.accurateBlockThreshold as the average block size. A skewed partition whose rows are spread over many map tasks has no individual block above that threshold, so its size is averaged away, adaptive query execution sees min == median == max and the skew is left unhandled. spark.shuffle.accurateBlockSkewedFactor, added by SPARK-36967, keeps the sizes of skewed blocks accurate, but it shipped disabled because enabling it sorted the block sizes in every map task. Only two order statistics are needed, so they are now selected in O(numPartitions) instead of sorting, and the feature is enabled by default with the value its documentation already recommended.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Two default-enabled P1 regressions are reproducible at the current head: cutoff ties bypass the configured skewed-block cap, and midpoint-pivot selection becomes quadratic on valid partition-size distributions. Please see the inline comments for concrete reproductions and suggested fixes.
| val sizes = uncompressedSizes.clone() | ||
| val medianSize: Long = Utils.medianInPlace(sizes) | ||
| val smallestAccurateSize = | ||
| Utils.nthSmallest(sizes, totalNumBlocks - maxAccurateSkewedBlockNumber) |
There was a problem hiding this comment.
[P1] Enforce the skewed-block cap when the cutoff has ties
nthSmallest returns the 100th-largest block size, but the classification below records every block with size >= threshold. This means ties at the cutoff bypass spark.shuffle.maxAccurateSkewedBlockNumber.
For example, with the proposed defaults and 50,000 reduce partitions, consider 25,001 blocks of 1 KiB and 24,999 blocks of 8 KiB. The median is 1 KiB, the 100th-largest size is 8 KiB, and the threshold is 8 KiB, so all 24,999 larger blocks are recorded instead of at most 100. writeExternal writes at least five bytes per recorded entry, so a 10,000-map stage produces 1,249,950,000 bytes of uncompressed map-status entries, before the much larger driver-side collection overhead. The base branch's disabled default does not record any of these blocks.
Since this PR enables this path by default, please enforce the actual cardinality bound, including deterministic handling of ties, while preserving mandatory accurate recording above spark.shuffle.accurateBlockThreshold. Please also add a regression test containing more than 100 equal-sized blocks at the cutoff.
There was a problem hiding this comment.
You're right, and the reproduction is exact. Fixed in 2ced43b.
The cutoff size alone was never a cardinality bound, only a size bound. The blocks strictly larger than the cutoff are now counted up front -- the selection leaves them all in the top-K window, so that count is at most the configured maximum -- and the remainder of the budget admits blocks of exactly the cutoff size, in block index order, so the map status stays deterministic. No strictly larger block is ever dropped in favour of a tied one, since only ties are rationed. Blocks at or above spark.shuffle.accurateBlockThreshold are still recorded unconditionally, so the mandatory path is unchanged.
The regression test is your example verbatim: 50000 partitions, 25001 blocks of 1 KiB and 24999 of 8 KiB, stock defaults. Without the bound it fails with 1024 did not equal 4600 -- all 24999 tied blocks get recorded and avgSize collapses to the size of the small blocks, which is a neat second symptom of the same bug.
| while (low < high) { | ||
| // The middle element keeps already sorted and reverse sorted inputs, which are both common | ||
| // for shuffle block sizes, away from the quadratic worst case. | ||
| val pivot = sizes(low + (high - low) / 2) |
There was a problem hiding this comment.
[P1] Bound the worst-case cost of default-enabled quickselect
Always selecting the middle element makes the algorithm quadratic for a valid organ-pipe reducer-size distribution: Array.tabulate(n)(i => math.min(i, n - 1 - i).toLong). Replaying the exact selections made by HighlyCompressedMapStatus produces 805,003 comparisons at 2,048 partitions, 3,174,774 at 4,096, 12,682,753 at 8,192, and 50,455,932 at 16,384. Doubling the number of partitions approximately quadruples the work.
This becomes a production regression because the PR changes spark.shuffle.accurateBlockSkewedFactor from -1.0 to 5.0, making every map task in a qualifying shuffle execute these selections; the base default performs no selection. At 16,384 partitions, a 10,000-map stage would require more than 504 billion comparisons just to construct map statuses.
Please use a worst-case-bounded selection algorithm, or an introspective/randomized pivot with a safe fallback, and add an organ-pipe input to the tests and benchmark.
There was a problem hiding this comment.
Confirmed, and thanks for the organ pipe input specifically. Fixed in 2ced43b.
nthSmallest is now an introselect: median-of-three pivot, and it sorts the range still under consideration once the ranges have stopped shrinking for 2 * log2(n) rounds. That bounds the worst case at O(n log n) regardless of the input, rather than relying on the pivot choice being lucky.
One thing I want to be upfront about: I did not reproduce your comparison counts or instrument the post-fix ones. What I added is a time bound on the input you named -- six selections over a 2^20 organ pipe under an explicit 60 second limit, which currently run in 100ms total. Extrapolating your table (50,455,932 at 16384, roughly quadrupling per doubling) puts a quadratic selection at more than 10^11 comparisons there, so the limit is generous enough not to be flaky while still failing decisively on a regression. If you would rather have an actual comparison count asserted against a ceiling, I can add one, though it means an instrumented copy of the loop in test code.
I also added a test asserting that nthSmallest leaves the array partitioned around the returned element, since HighlyCompressedMapStatus now counts the sizes above the cutoff by scanning only the tail of the array and so depends on that ordering, not just on the returned value.
The benchmark now reports the organ pipe distribution alongside the skewed one. At 50000 partitions it is 530us per map task against 310us for the ordinary case -- a 1.7x gap, not a blowup.
…lection cost Address two problems with enabling the accurate recording of skewed block sizes by default. The number of accurately recorded blocks was only bounded by the size of the cutoff block, so blocks tied at that size were all recorded and could exceed spark.shuffle.maxAccurateSkewedBlockNumber by orders of magnitude. With 50000 reduce partitions, 25001 blocks of 1 KiB and 24999 blocks of 8 KiB, all 24999 larger blocks were recorded instead of at most 100. The blocks strictly larger than the cutoff are now counted up front, and the rest of the budget admits the tied blocks in block index order, which keeps the map status deterministic. Blocks at or above spark.shuffle.accurateBlockThreshold are still always recorded, as they were before. Utils.nthSmallest picked the middle element as the pivot, which is quadratic for an organ pipe distribution of block sizes. It is now an introselect: the pivot is the median of the first, middle and last element, and the range still under consideration is sorted once the ranges have stopped shrinking quickly enough, which bounds the worst case at O(numPartitions * log(numPartitions)). Both cases are covered by new tests, the ordering that the block count relies on is now asserted directly, and the benchmark also reports the organ pipe distribution.
2 participants
What changes were proposed in this pull request?
Once a shuffle has more than
spark.shuffle.minNumPartitionsToHighlyCompress(2000) partitions, map statuses switch toHighlyCompressedMapStatus, which keeps exact sizes only for blocks abovespark.shuffle.accurateBlockThreshold(100MB) and reports the average block size for everything else.MapOutputTracker.getStatisticssums those per reduce id, so AQE ends up looking at a flat distribution.spark.shuffle.accurateBlockSkewedFactor, added by SPARK-36967, keeps skewed block sizes accurate, but it ships disabled (-1.0) even though its own documentation recommends setting it to the same value asspark.sql.adaptive.skewJoin.skewedPartitionFactor. Enabling it fully sorted the block sizes in every map task, which is presumably why it was left off.This PR:
HighlyCompressedMapStatus.applywith selection of the only two order statistics it needs (the median, and the K-th largest size), using a new in-placeUtils.nthSmallestplusUtils.medianInPlace. This is O(numPartitions) on average rather than O(numPartitions log numPartitions), with the same result.nthSmallestis an introselect: the pivot is the median of the first, middle and last element, and it sorts the range still under consideration once the ranges have stopped shrinking quickly enough, which bounds the worst case at O(numPartitions log numPartitions) even for adversarial inputs such as an organ pipe distribution;spark.shuffle.maxAccurateSkewedBlockNumberas an actual cardinality bound. Previously only the size of the cutoff block bounded the recording, so blocks tied at that size were all recorded. The blocks strictly larger than the cutoff are now counted up front, and the remainder of the budget admits tied blocks in block index order, which keeps the map status deterministic. Blocks at or abovespark.shuffle.accurateBlockThresholdare still always recorded, as they were before;spark.shuffle.accurateBlockSkewedFactorfrom-1.0to5.0, matchingspark.sql.adaptive.skewJoin.skewedPartitionFactor.-1.0still disables the feature;Why are the changes needed?
Skew join optimization silently stops working above 2000 shuffle partitions. In the reporter's case a 263GB partition spread over 10335 mappers is only ~25MB per block, far below the 100MB accurate threshold, so every block is reported as the average and AQE logs
median size == max size == min sizeand finds no skew. Loweringspark.shuffle.accurateBlockThresholddoes not help, because what matters is the per-block size rather than the per-partition total.Reproduced on master with a join where 95% of the rows land on one key, varying only
spark.sql.shuffle.partitions:spark.sql.shuffle.partitionsSortMergeJoin(skew=true)spark.shuffle.accurateBlockSkewedFactor=5.0spark.shuffle.minNumPartitionsToHighlyCompress=100000Either knob restores detection on its own, which locates the problem in the map status compression rather than in AQE.
The cost of the change, measured with the benchmark added here (per map task, Apple M4, JDK 21):
For comparison, the sorting implementation this PR replaces measured 28.8us, 140.7us and 1562.0us on the same machine, so at 50000 partitions the accurate path is now about 5x cheaper than sorting would have been. The organ pipe column is the worst case for the selection, where the introselect fallback kicks in; it stays within 1.7x of the ordinary case and does not grow quadratically.
On the driver side each map status carries at most
spark.shuffle.maxAccurateSkewedBlockNumber(100) extra entries of an int and a byte, so about 500 bytes per map task, or ~5MB for a 10000 map task stage.Does this PR introduce any user-facing change?
Yes. Shuffles with more than
spark.shuffle.minNumPartitionsToHighlyCompresspartitions now report skewed block sizes accurately by default, so AQE can detect and split skewed partitions where it previously could not. Queries affected by this will change plans (skew joins get split) and should run faster. Settingspark.shuffle.accurateBlockSkewedFactor=-1.0restores the old behavior.How was this patch tested?
New tests in
MapStatusSuite:spark.shuffle.accurateBlockThresholdis recorded accurately under default configuration, and the average is not inflated by it. Fails on master with88733 did not equal 1500.spark.shuffle.maxAccurateSkewedBlockNumberblocks are recorded and the rest report the average. Fails without the cardinality bound with1024 did not equal 4600, because all 24999 tied blocks were recorded and the average collapsed to the size of the small blocks.New tests in
UtilsSuite:Utils.nthSmallestagrees with a full sort over sorted, reverse sorted, constant, organ pipe, duplicate-heavy and random inputs, and its argument validation.Utils.nthSmallestleaves the array partitioned around the returned element.HighlyCompressedMapStatuscounts the sizes above the cutoff by scanning only the tail of the array, so it depends on this and not just on the returned value.Utils.medianInPlaceagrees withUtils.median.HighlyCompressedMapStatusBenchmarkis new, with results generated at 2048, 10000 and 50000 partitions for both distributions.Existing suites:
MapStatusSuite,MapOutputTrackerSuite,ShuffleSuite,UtilsSuite,org.apache.spark.shuffle.sort.*andAdaptiveQueryExecSuite(129 tests) all pass.End to end, with a locally built assembly and no configuration beyond the SQL level skew join settings, the table above becomes yes / yes / yes at 1999, 2001 and 4000 partitions.
I did not add an AQE level test for the >2000 partition case:
spark.shuffle.minNumPartitionsToHighlyCompressis read through a lazy val on the JVM'sSparkEnv, so it cannot be lowered from within a suite, and running a query with 2001 shuffle partitions takes about a minute. SPARK-36967 tested this feature atMapStatusSuitelevel for the same reason. Happy to add one if reviewers would rather pay the runtime.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code